from sqlmodel import SQLModel, Field, Relationship, create_engine, Session, select


# 一个实例就是一个类，一次记录就是一个实例，字段代表属性
# 多对多关系的额外字段（如成绩）应存储在关联表中，而非主实体表
# 先提交主实体，再设置关联表字段，确保外键关系已建立
# 关联表（联结表，中间表）
class Enrollment(SQLModel, table=True):
    student_id: int = Field(foreign_key="students.id", primary_key=True)
    course_id: int = Field(foreign_key="courses.id", primary_key=True)
    grade: str | None = None  # 可选字段：表示学生在该课程的成绩


# 主实体：学生
class Student(SQLModel, table=True):
    __tablename__ = "students"

    id: int = Field(primary_key=True, default=None)
    name: str

    # 定义多对多关系：一个学生可以选修多门课程
    courses: list["Course"] = Relationship(
        back_populates="students", link_model=Enrollment  # 指定关联表
    )


# 主实体：课程
class Course(SQLModel, table=True):
    __tablename__ = "courses"

    id: int = Field(primary_key=True, default=None)
    title: str

    # 定义多对多关系：一门课程可以被多个学生选修
    students: list["Student"] = Relationship(
        back_populates="courses", link_model=Enrollment  # 指定关联表
    )


# 配置并创建数据库
sqlite_filepath = "school.db"
sqlite_url = f"sqlite:///{sqlite_filepath}"
engine = create_engine(sqlite_url)


def create_db_and_tables():
    SQLModel.metadata.create_all(engine)


def insert_data():
    with Session(engine) as session:
        alice = Student(name="Alice")
        bob = Student(name="Bob")
        math = Course(title="Mathematics")
        physics = Course(title="Physics")

        alice.courses.append(math)
        alice.courses.append(physics)
        bob.courses.append(math)

        session.add(math)
        session.add(bob)
        session.commit()
        session.refresh(alice)
        session.refresh(bob)

        for course in alice.courses:
            if course.title == "Mathematics":
                enrollment = session.exec(
                    select(Enrollment).where(
                        Enrollment.student_id == alice.id,
                        Enrollment.course_id == course.id,
                    )
                ).one()
                enrollment.grade = "A"
        session.commit()
        print("数据插入成功！")
        return alice, bob, math, physics


def query_data():
    with Session(engine) as session:
        students = session.exec(select(Student)).all()
        for student in students:
            print(f"\n学生:{student.name}")
            print("选修课程:")
            for course in student.courses:
                enrollment = session.exec(
                    select(Enrollment).where(
                        Enrollment.student_id == student.id,
                        Enrollment.course_id == course.id,
                    )
                ).one()
                grade = enrollment.grade if enrollment.grade else "未评分"
                print(f" - {course.title}(成绩: {grade})")

        courses = session.exec(select(Course))
        for course in courses:
            print(f"\n学生:{student.name}")
            print("选修课程:")
            for student in course.students:
                # 查询关联表中的成绩
                enrollment = session.exec(
                    select(Enrollment).where(
                        Enrollment.student_id == student.id,
                        Enrollment.course_id == course.id,
                    )
                ).one()
                grade = enrollment.grade if enrollment.grade else "未评分"
                print(f"  - {student.name} (成绩: {grade})")


if __name__ == "__main__":
    # create_db_and_tables()
    # insert_data()
    query_data()
